#include "comm/znx_rbtree.h"

typedef struct {
    znx_rbtree_node_t   tree_node;
    int                 key;
    int                 val;
} znx_test_item_t;

static int
znx_test_item_cmp(znx_rbtree_node_t *a, znx_rbtree_node_t *b) {
    znx_test_item_t *anode, *bnode;
    anode = (znx_test_item_t *)a;
    bnode = (znx_test_item_t *)b;

    if (anode->key == bnode->key) {
        return 0;
    }

    if (anode->key > bnode->key) {
        return 1;
    }

    return -1;
}

void
test_rbtree() {
    znx_rbtree_t tree;
    znx_rbtree_node_t sentinel;
    znx_rbtree_init(&tree, &sentinel, znx_test_item_cmp);

    znx_test_item_t items[20];

    for (int i = 0; i < 20; i++) {
        items[i].key = 2*i;
        znx_rbtree_insert(&tree, &items[i].tree_node);
    }

    for (int i = 0; i < 20; i++) {
        znx_test_item_t target = {.key = 2*i};
        znx_rbtree_node_t *res = znx_rbtree_lookup(&tree, &target.tree_node);
        assert(res != NULL);

        znx_test_item_t *t = (znx_test_item_t *)res;
        assert(t->key == 2 * i);
    }

    znx_rbtree_node_t *tmp;
    for (int i = 0; i < 20; i++) {
        znx_test_item_t target = {.key = 2*i};
        tmp = znx_rbtree_lookup(&tree, &target.tree_node);
        znx_rbtree_delete(&tree, tmp);
    }

    assert(znx_rbtree_empty(&tree));

    for (int i = 0; i < 20; i++) {
        items[i].key = 2*i;
        znx_rbtree_insert(&tree, &items[i].tree_node);
    }

    for (int i = 0; i < 20; i++) {
        znx_test_item_t target = {.key = 0};
        znx_rbtree_node_t *res = znx_rbtree_ge(&tree, &target.tree_node);
        assert(res != NULL);

        znx_test_item_t *t = (znx_test_item_t *)res;
        assert(t->key == 2 * i);

        znx_rbtree_delete(&tree, res);
    }

    assert(znx_rbtree_empty(&tree));

    for (int i = 0; i < 20; i++) {
        items[i].key = 2*i;
        znx_rbtree_insert(&tree, &items[i].tree_node);
    }

    for (int i = 19; i >= 0; i--) {
        znx_test_item_t target = {.key = 20*2};
        znx_rbtree_node_t *res = znx_rbtree_le(&tree, &target.tree_node);
        assert(res != NULL);

        znx_test_item_t *t = (znx_test_item_t *)res;
        assert(t->key == 2 * i);

        znx_rbtree_delete(&tree, res);
    }
    
    assert(znx_rbtree_empty(&tree));

    // other
    for (int i = 0; i < 20; i++) {
        items[i].key = 2*i;
        znx_rbtree_insert(&tree, &items[i].tree_node);
    }

    znx_test_item_t target = {.key = -1};
    assert(znx_rbtree_lookup(&tree, &target.tree_node) == NULL);

    assert(znx_rbtree_le(&tree, &target.tree_node) == NULL);

    target.key = 20 * 2 + 1;
    assert(znx_rbtree_lookup(&tree, &target.tree_node) == NULL);

    assert(znx_rbtree_ge(&tree, &target.tree_node) == NULL);
}

void test_rbtree_same_key() {
    znx_rbtree_t tree;
    znx_rbtree_node_t sentinel;
    znx_rbtree_init(&tree, &sentinel, znx_test_item_cmp);

    znx_test_item_t items[20];
    for (int i = 0; i < 10; i++) {
        items[i].key = 1;
        items[i].val = i;
        znx_rbtree_insert(&tree, &items[i].tree_node);
    }

    for (int i = 10; i < 20; i++) {
        items[i].key = 2;
        items[i].val = i;
        znx_rbtree_insert(&tree, &items[i].tree_node);
    }

    for (int i = 0; i < 20; i++) {
        znx_rbtree_node_t *node = znx_rbtree_min(tree.root, tree.sentinel);
        assert(node != NULL);

        znx_test_item_t *item = (znx_test_item_t *)node;
        if (i < 10) {
            assert(item->key == 1);
            DD("get item->key: %d, item->value: %d",item->key, item->val);
        } else {
            assert(item->key == 2);
            DD("get item->key: %d, item->value: %d", item->key, item->val);
        }

        znx_rbtree_delete(&tree, node);
    }

    assert(znx_rbtree_empty(&tree));
}

int main() {
    test_rbtree();
    test_rbtree_same_key();
    return 0;
}